{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import tempfile\n",
    "import numpy as np\n",
    "import networkx as nx\n",
    "from rdkit.Chem import AllChem\n",
    "from collections import Counter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# functions to reformat the google seq2seq beam files\n",
    "\n",
    "def decode_single_example(predicted_ids, parent_ids, scores, vocab,\n",
    "                          num_top_candidates, calc_num_completed=False):\n",
    "  \"\"\"\n",
    "  Decodes a numpy array of predicted ids, parent_ids, scores, each with\n",
    "  dimensions seq_length x k (# of beams), into a python list of tuples.\n",
    "  Each tuple has the form ((chars), path log prob, path prob) and represents\n",
    "  a possible beam candidate. Adapted\n",
    "  from https://gist.github.com/MInner/5204e649f6a7b0541b232f1f0f9fc8ba. The\n",
    "  number of beam candidates will be <= N\n",
    "  :param predicted_ids: numpy array of predicted ids, dim seq_length x k\n",
    "  :param parent_ids:  numpy array of predicted ids, dim seq_length x k\n",
    "  :param scores:  numpy array of predicted ids, dim seq_length x k\n",
    "  :param vocab: a list where the index corresponds to the char in the vocab\n",
    "  :param num_top_candidates: the number of top candidates (N) to keep for each\n",
    "  example, where N <= k (beam width used in seq2seq decoder). The actual\n",
    "  number of candidates generated will always be less than or equal to N.\n",
    "  :param calc_num_completed: if True, also returns the number of completed\n",
    "  candidates (ie candidates that\n",
    "  :return: a list of tuples, each of the form\n",
    "  ((chars), path log prob, path prob). The prob is weighted by the number of\n",
    "  candidates in the list. Eg if the list contains 1 candidate, then the prob\n",
    "  of that candidate will be 1.0\n",
    "  \"\"\"\n",
    "  k = predicted_ids.shape[1]\n",
    "  if num_top_candidates > k:\n",
    "    raise ValueError(\"num_top_candidates cannot be greater than k\")\n",
    "\n",
    "  def _add_graph_level(graph, level, parent_ids, names, scores):\n",
    "    \"\"\"Adds a level to the passed graph\"\"\"\n",
    "    for i, parent_id in enumerate(parent_ids):\n",
    "      new_node = (level, i)\n",
    "      parent_node = (level - 1, parent_id)\n",
    "      graph.add_node(new_node)\n",
    "      graph.node[new_node][\"name\"] = names[i]\n",
    "      graph.node[new_node][\"score\"] = str(scores[i])\n",
    "      graph.node[new_node][\"size\"] = 100\n",
    "      # Add an edge to the parent\n",
    "      graph.add_edge(parent_node, new_node)\n",
    "\n",
    "  def create_graph(predicted_ids, parent_ids, scores, vocab=None):\n",
    "    def get_node_name(pred):\n",
    "      return vocab[pred] if vocab else str(pred)\n",
    "\n",
    "    seq_length = predicted_ids.shape[0]\n",
    "    graph = nx.DiGraph()\n",
    "    for level in range(seq_length):\n",
    "      names = [get_node_name(pred) for pred in predicted_ids[level]]\n",
    "      _add_graph_level(graph, level + 1, parent_ids[level], names,\n",
    "                       scores[level])\n",
    "    graph.node[(0, 0)][\"name\"] = \"START\"\n",
    "    return graph\n",
    "\n",
    "  def get_path_to_root(graph, node):\n",
    "    p = graph.predecessors(node)\n",
    "    assert len(p) <= 1\n",
    "    self_seq = [graph.node[node]['name'].split('\\t')[0]]\n",
    "    if len(p) == 0:\n",
    "      return self_seq\n",
    "    else:\n",
    "      return self_seq + get_path_to_root(graph, p[0])\n",
    "\n",
    "  def _tree_node_predecessor(pos):\n",
    "    return graph.node[graph.predecessors(pos)[0]]\n",
    "\n",
    "  graph = create_graph(predicted_ids, parent_ids, scores, vocab)\n",
    "\n",
    "  pred_end_node_names = {pos for pos, d in graph.node.items()\n",
    "                         if d['name'] == 'SEQUENCE_END'\n",
    "                         and len(graph.predecessors(pos)) > 0\n",
    "                         and _tree_node_predecessor(pos)[\n",
    "                           'name'] != 'SEQUENCE_END'}\n",
    "\n",
    "  # if no suitable end nodes, occurs when all candidates have been cutoff\n",
    "  # prematurely due to exceeding the max seq2seq decode length. We create a\n",
    "  # dummy output that is an invalid molecule SMILES, with nan for both log\n",
    "  # probs and scores\n",
    "  if len(pred_end_node_names) == 0:\n",
    "    if calc_num_completed:\n",
    "      return [((\"-\"), float(\"nan\"), float(\"nan\"))], 0\n",
    "\n",
    "    return [((\"-\"), float(\"nan\"), float(\"nan\"))]\n",
    "\n",
    "  result = [(tuple(get_path_to_root(graph, pos)[1:-1][::-1]),\n",
    "             float(graph.node[pos]['score']))\n",
    "            for pos in pred_end_node_names]\n",
    "\n",
    "  filtered_result = filter(lambda x: 'SEQUENCE_END' not in x[0], result)\n",
    "\n",
    "  s_result = sorted(filtered_result, key=lambda x: x[1], reverse=True)\n",
    "\n",
    "  # take the top N candidates, sorted by decreasing log prob\n",
    "  top_N_result = s_result[:num_top_candidates]\n",
    "\n",
    "  nn_probs = np.exp(np.array(list(zip(*top_N_result))[1]))\n",
    "  probs = nn_probs / np.sum(nn_probs)\n",
    "  result_w_prob = [(path, score, prob) for (path, score), prob in\n",
    "                   zip(s_result, probs)]  #NB: the length of the\n",
    "  # result_w_prob will be <= N\n",
    "\n",
    "  if calc_num_completed:\n",
    "    return result_w_prob, len(pred_end_node_names)\n",
    "\n",
    "  return result_w_prob\n",
    "\n",
    "def create_beam_source_target_parallel_txt_files(beam_path, vocab_path,\n",
    "                                                 output_pred_path, input_target_path,\n",
    "                                                 output_target_path,\n",
    "                                                 input_source_path,\n",
    "                                                 output_source_path,\n",
    "                                                 num_top_candidates):\n",
    "  \"\"\"\n",
    "  Converts the beam.npz data from the google seq2seq inference with a beam\n",
    "  decoding width of k into a pred file containing the SMILES of the top k\n",
    "  candidates, with each character separated by a space, for each example in\n",
    "  order. NB: currently there is no guarantee that each example will have\n",
    "  exactly k candidates. Also reformats the original target and source files so\n",
    "  that each candidate in the beam source file corresponds to the correct\n",
    "  target and source in parallel txt format.\n",
    "  :param beam_path: path to the beam.npz file\n",
    "  :param vocab_path: path to the vocab file\n",
    "  :param output_pred_path: path to the output pred file\n",
    "  :param input_target_path: path to original target file\n",
    "  :param output_target_path:  path to reformatted target file\n",
    "  :param input_source_path: path to original source file\n",
    "  :param output_source_path:  path to reformatted source file\n",
    "  :param num_top_candidates: the number of top candidates (N) to keep for each\n",
    "  example, where N <= k (beam width used in seq2seq decoder)\n",
    "  :return:\n",
    "  \"\"\"\n",
    "\n",
    "  def _read_vocab_from_txt(vocab_path):\n",
    "    \"\"\"\n",
    "    Reads the vocabulary file from the vocab path into a list, and also adds\n",
    "    the special tokens from the google seq2seq model\n",
    "    :param vocab_path:\n",
    "    :return:  a list where the index corresponds to the char\n",
    "    \"\"\"\n",
    "    with open(vocab_path, \"r\") as infile:\n",
    "      vocab = infile.readlines()\n",
    "    vocab = [_.strip() for _ in vocab]\n",
    "    vocab += [\"UNK\", \"SEQUENCE_START\", \"SEQUENCE_END\"]\n",
    "\n",
    "    return vocab\n",
    "\n",
    "  vocab = _read_vocab_from_txt(vocab_path)\n",
    "  beam_data = np.load(beam_path)\n",
    "\n",
    "  all_predicted_ids = beam_data[\"predicted_ids\"]\n",
    "  all_parent_ids = beam_data[\"beam_parent_ids\"]\n",
    "  all_scores = beam_data[\"scores\"]\n",
    "\n",
    "  num_examples = all_predicted_ids.shape[0]\n",
    "\n",
    "  with open(output_pred_path, \"w\") as out_pred, \\\n",
    "      open(input_target_path, \"r\") as in_target, \\\n",
    "      open(output_target_path, \"w\") as out_target, \\\n",
    "    open(input_source_path, \"r\") as in_source, \\\n",
    "    open(output_source_path, \"w\") as out_source:\n",
    "\n",
    "    for i in range(num_examples):\n",
    "      decoded = decode_single_example(all_predicted_ids[i],\n",
    "                                      all_parent_ids[i], all_scores[i], vocab,\n",
    "                                      num_top_candidates)\n",
    "\n",
    "      target = in_target.readline() # target for this example, includes \\n\n",
    "      source = in_source.readline() # source for this example, includes \\n\n",
    "\n",
    "      for candidate in decoded:\n",
    "        smiles_char_string = \" \".join(candidate[0])  # SMILES string with each\n",
    "        # char separated by a space\n",
    "        out_pred.write(smiles_char_string + \"\\n\")\n",
    "        out_target.write(target)\n",
    "        out_source.write(source)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# functions to evaluate the predictions\n",
    "\n",
    "def single_exact_match_score(prediction, ground_truth, canonicalize=False):\n",
    "  \"\"\"\n",
    "  Calculates the exact match score for a single prediction and ground truth\n",
    "  pair, where they represent the reactant parts of a reaction SMILES string.\n",
    "  Uses InChI comparison if canonicalize\n",
    "  :param prediction:  a string of tokens either separated or not by space\n",
    "  delimiters\n",
    "  :param ground_truth:  a string of tokens separated or not by space delimiters\n",
    "  :param canonicalize:  if True, canonicalizes the prediction and ground\n",
    "  truth reactant SMILES prior to the exact match comparison if both the\n",
    "  prediction and ground truth are valid reactant SMILES.\n",
    "  :return:  1.0 if exact match, 0.0 otherwise\n",
    "  \"\"\"\n",
    "  if canonicalize:\n",
    "    prediction_smiles = \"\".join(prediction.split())\n",
    "    ground_truth_smiles = \"\".join(ground_truth.split())\n",
    "    pred_mol_obj = AllChem.MolFromSmiles(prediction_smiles)\n",
    "    ground_truth_mol_obj = AllChem.MolFromSmiles(ground_truth_smiles)\n",
    "    if pred_mol_obj and ground_truth_mol_obj: # if either is a None obj,\n",
    "  # then that corresponding SMILES could not be parsed\n",
    "      prediction = AllChem.MolToInchi(pred_mol_obj)\n",
    "      ground_truth = AllChem.MolToInchi(ground_truth_mol_obj)\n",
    "\n",
    "  return float(prediction == ground_truth)\n",
    "\n",
    "def check_valid_smiles(input_smiles):\n",
    "  \"\"\"\n",
    "  Returns True if the input species SMILES is overall valid gramatically\n",
    "  :param input_smiles: a species smiles to check (can be multiple species\n",
    "  separated by .)\n",
    "  :return: True if input species SMILES is valid, False if otherwise\n",
    "  \"\"\"\n",
    "  m = AllChem.MolFromSmiles(input_smiles)\n",
    "  if m:\n",
    "    return True\n",
    "  else:\n",
    "    return False\n",
    "\n",
    "def nested_list_to_1d_array(nested_list):\n",
    "    \"\"\"\n",
    "    Converts a list of lists into a 1d array. Guarantees that the resulting\n",
    "    array will be 1d, even if the nested list contains lists of equal length.\n",
    "    Addresses the inconsistent behaviour of np.array array creation, since a\n",
    "    list of lists of unequal length results in a 1d array, while a list of\n",
    "    lists of equal length results in a nd array.\n",
    "    :param nested_list: a list of lists\n",
    "    :return: a 1d numpy array of Dim (num of nested lists,)\n",
    "    \"\"\"\n",
    "    len_list = len(nested_list) # num of nested lists\n",
    "    result = np.empty(len_list, dtype=np.object)\n",
    "    result[:] = nested_list\n",
    "\n",
    "    return result\n",
    "\n",
    "def evaluate_beam_by_rxn_class(beam_prediction_path,\n",
    "                                    beam_ground_truth_path,\n",
    "                               beam_source_path, result_path):\n",
    "  \"\"\"\n",
    "  Evaluates a prediction file and ground truth file pair in parallel text\n",
    "  format by reaction class, and writes the metrics and correct/incorrect\n",
    "  predictions to a results file. The reaction\n",
    "  class information is contained in the test source file that was used as\n",
    "  input in the inference step to obtain the prediction file.\n",
    "  :param beam_prediction_path: path to the beam prediction file, resulting\n",
    "  from create_beam_source_target_parallel_txt_files\n",
    "  :param beam_ground_truth_path: path to the reformatted target file,\n",
    "  resulting from create_beam_source_target_parallel_txt_files\n",
    "  :param beam_source_path:  path to the reformatted source file, resulting from\n",
    "  create_beam_source_target_parallel_txt_files\n",
    "  :param result_path: path to output file\n",
    "  :return:\n",
    "  \"\"\"\n",
    "  def _read_list_from_txt(input_path):\n",
    "    \"\"\"\n",
    "    Reads a text file containing lines of chars into a list of list of the\n",
    "    elements\n",
    "    :param input_path: path to input file\n",
    "    :return:  a list of list of the elements\n",
    "    \"\"\"\n",
    "    with open(input_path, 'r') as infile:\n",
    "      list_of_lists = []\n",
    "      for line in infile:\n",
    "        list_of_lists.append([i for i in line.rstrip().split()])\n",
    "    return list_of_lists\n",
    "\n",
    "  def _evaluate_array(prediction_array, ground_truth_array,\n",
    "                      test_source_array, result_path, reaction_class=None):\n",
    "    \"\"\"\n",
    "    Evaluates a prediction array and ground truth array pair in parallel text\n",
    "    format, and appends the metrics and correct/incorrect predictions to a\n",
    "    results file\n",
    "    :param prediction_array: 1d numpy array of lists of elements\n",
    "    :param ground_truth_array:  1d numpy array of lists of elements\n",
    "    :param result_path: path to output file\n",
    "    :param reaction_class:  a reaction class character\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    # initialize global metrics\n",
    "    total_em_score = 0  # if any of the beam candidates of an example is an\n",
    "    # exact match, we score that example as correct. Theoretically, the max\n",
    "    # possible total_em_score is num_examples.\n",
    "    total_invalid_smiles = 0\n",
    "    total_valid_smiles = 0\n",
    "\n",
    "    # lists will eventually have len num_examples, where each entry\n",
    "    # corresponds to the num of em predictions for each example\n",
    "    num_em_per_example_list = []\n",
    "\n",
    "    num_examples = 0\n",
    "\n",
    "    # initialize lists for correct/incorrect predictions\n",
    "    correct = []  # preds only (ie only reactants), not whole rxn smiles\n",
    "    incorrect = []\n",
    "\n",
    "    num_candidates = prediction_array.shape[0]  # each target can have\n",
    "    # multiple candidates due to beam search\n",
    "\n",
    "    current_target = \" \".join(ground_truth_array[0])  # initialize current\n",
    "    # target to the first target\n",
    "\n",
    "    # initialize current target (not candidate) metrics\n",
    "    current_em = 0\n",
    "    for i in range(num_candidates):\n",
    "      prediction = \" \".join(prediction_array[i])  # to get into suitable\n",
    "      # format for the single metric calculations, which take in a string of\n",
    "      # elements separated by spaces\n",
    "      truth = \" \".join(ground_truth_array[i])\n",
    "\n",
    "      # count valid and invalid prediction smiles\n",
    "      if check_valid_smiles(\"\".join(prediction_array[i])):\n",
    "        total_valid_smiles += 1\n",
    "      else:\n",
    "        total_invalid_smiles += 1\n",
    "\n",
    "      # when the current candidate is from a different example/target,\n",
    "      # update current target and metrics for the previous example/target\n",
    "      if truth != current_target:\n",
    "        # update global metrics with current target metrics\n",
    "        total_em_score += current_em\n",
    "        num_em_per_example_list.append(current_em)\n",
    "        # reinitialize current target metrics\n",
    "        current_em = 0\n",
    "        # update the current target\n",
    "        current_target = truth\n",
    "        num_examples += 1\n",
    "\n",
    "      # calculate single example metrics and update\n",
    "      single_em = single_exact_match_score(prediction, truth, canonicalize=True)\n",
    "      # have this as a separate variable to be used for indicate\n",
    "      # correct/incorrect pred\n",
    "      current_em += single_em\n",
    "\n",
    "      # format the product source SMILES without the reaction char, so that\n",
    "      # later we can formulate the reaction SMILES. Also format the pred and\n",
    "      # truth SMILES\n",
    "      product_smiles = \"\".join(test_source_array[i][1:])\n",
    "      pred_smiles = \"\".join(prediction_array[i])\n",
    "      truth_smiles = \"\".join(ground_truth_array[i])\n",
    "\n",
    "      # update correct or incorrect list\n",
    "      if single_em:\n",
    "        correct.append((pred_smiles,\n",
    "                        truth_smiles,\n",
    "                        product_smiles))\n",
    "      else:\n",
    "        incorrect.append((pred_smiles,\n",
    "                          truth_smiles,\n",
    "                          product_smiles))\n",
    "\n",
    "      # once we reach the last candidate, need to update the metrics for the\n",
    "      # this example/target. #TODO(Bowen): pretty hacky. A lot of repeated\n",
    "      # code from the earlier check whether truth != current_target\n",
    "      if i == range(num_candidates)[-1]:\n",
    "        # update global metrics with current target metrics\n",
    "        total_em_score += current_em\n",
    "        num_em_per_example_list.append(current_em)\n",
    "        num_examples += 1\n",
    "\n",
    "    # calculate final metrics\n",
    "    num_examples_with_em = sum(x > 0 for x in num_em_per_example_list)\n",
    "    average_em = num_examples_with_em / num_examples  # this is the top k accuracy\n",
    "\n",
    "    # write result file\n",
    "    with open(result_path, \"a\") as result_file:\n",
    "      if reaction_class:\n",
    "        result_file.write(\"###### {} ######\".format(reaction_class) + \"\\n\")\n",
    "      else:\n",
    "        result_file.write(\"###### OVERALL ######\" + \"\\n\")\n",
    "      result_file.write(\"METRICS\" + \"\\n\")\n",
    "      result_file.write(\"Number of examples: {}\".format(num_examples) + \"\\n\")\n",
    "      result_file.write(\n",
    "        \"Number of candidates: {}\".format(num_candidates) + \"\\n\")\n",
    "      result_file.write(\"Number of examples with exact matches: {}\".format((\n",
    "        num_examples_with_em)) + \"\\n\")\n",
    "      result_file.write(\"Top-k accuracy: {}\".format(average_em) + \"\\n\")\n",
    "\n",
    "      # record valid and invalid prediction smiles\n",
    "      result_file.write(\"Number of valid SMILES: {}\".format(\n",
    "        total_valid_smiles) + \"\\n\")\n",
    "      result_file.write(\"Number of invalid SMILES: {}\".format(\n",
    "        total_invalid_smiles) + \"\\n\")\n",
    "\n",
    "      result_file.write(\"\\n\")\n",
    "\n",
    "      result_file.write(\"INCORRECT PREDICTIONS\" + \"\\n\")\n",
    "      for pred, truth, prod in incorrect:\n",
    "        pred_rxn_smiles = pred + \">>\" + prod\n",
    "        truth_rxn_smiles = truth + \">>\" + prod\n",
    "        result_file.write(\"{} {}\".format(pred_rxn_smiles, truth_rxn_smiles) + \"\\n\")\n",
    "        \n",
    "      result_file.write(\"\\n\")\n",
    "\n",
    "      result_file.write(\"CORRECT PREDICTIONS\" + \"\\n\")\n",
    "      for pred, truth, prod in correct:\n",
    "        pred_rxn_smiles = pred + \">>\" + prod\n",
    "        truth_rxn_smiles = truth + \">>\" + prod\n",
    "        result_file.write(\"{} {}\".format(pred_rxn_smiles, truth_rxn_smiles) +\n",
    "                          \"\\n\")\n",
    "\n",
    "      result_file.write(\"\\n\\n\")\n",
    "\n",
    "  def _get_rxn_class_indices(rxn_class, test_source_array):\n",
    "    \"\"\"\n",
    "    Given a test source array of lists, in which each list contains a\n",
    "    reaction class character in the first entry, returns an array that has\n",
    "    entries of True where the corresponding list in the test source array\n",
    "    contains a reaction class character equal to rxn_class, and False otherwise\n",
    "    :param rxn_class: a reaction class character\n",
    "    :param test_source_array: a 1d numpy array of lists, where each list\n",
    "    contains a reaction class character in the first entry\n",
    "    :return:  a 1d numpy boolean array\n",
    "    \"\"\"\n",
    "    index_list = []\n",
    "    for list in test_source_array:\n",
    "      if list[0] == rxn_class:\n",
    "        index_list.append(True)\n",
    "      else:\n",
    "        index_list.append(False)\n",
    "    index_array = np.array(index_list)\n",
    "\n",
    "    assert index_array.shape == test_source_array.shape\n",
    "\n",
    "    return index_array\n",
    "\n",
    "  # read files into 1d numpy arrays. Use read_list_from_txt and then convert\n",
    "  # list of lists to array because the reaction class char is a str and\n",
    "  # cannot be converted into int\n",
    "  prediction_array = nested_list_to_1d_array(_read_list_from_txt(\n",
    "    beam_prediction_path))\n",
    "  ground_truth_array = nested_list_to_1d_array(_read_list_from_txt(\n",
    "    beam_ground_truth_path))\n",
    "  test_source_array = nested_list_to_1d_array(_read_list_from_txt(\n",
    "    beam_source_path))\n",
    "\n",
    "  # calculate metric and print correct/incorrect examples for overall array\n",
    "  _evaluate_array(prediction_array, ground_truth_array, test_source_array,\n",
    "                  result_path, reaction_class=None)\n",
    "\n",
    "  # find array of unique reaction classes from test source\n",
    "  rxn_classes = list(set([element[0] for element in test_source_array]))\n",
    "\n",
    "  # find indices for each reaction class and then create new class array.\n",
    "  # Calculate metrics and print correct/incorrect examples for each reaction\n",
    "  # class array\n",
    "  for rxn_class in rxn_classes:\n",
    "    rxn_class_indices = _get_rxn_class_indices(rxn_class, test_source_array)\n",
    "    class_prediction_array = prediction_array[rxn_class_indices]\n",
    "    class_ground_truth_array = ground_truth_array[rxn_class_indices]\n",
    "    class_test_source_array = test_source_array[rxn_class_indices]\n",
    "    _evaluate_array(class_prediction_array, class_ground_truth_array,\n",
    "                    class_test_source_array, result_path, rxn_class)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# evaluate the seq2seq predictions\n",
    "\n",
    "# define paths\n",
    "cwd = os.getcwd()\n",
    "prediction_dir = os.path.join(os.path.dirname(cwd), 'predictions')\n",
    "processed_data_dir = os.path.join(os.path.dirname(cwd), 'processed_data')\n",
    "\n",
    "vocab_path = os.path.join(processed_data_dir, 'vocab')\n",
    "target_path = os.path.join(processed_data_dir, 'test_targets')\n",
    "source_path = os.path.join(processed_data_dir, 'test_sources')\n",
    "\n",
    "# no beam (ie beam width = 1)\n",
    "pred_path = os.path.join(prediction_dir, 'predictions_44000_steps_beam_1_test.txt')\n",
    "result_path = os.path.join(cwd, 'beam_1_test_results')\n",
    "\n",
    "evaluate_beam_by_rxn_class(pred_path, target_path, source_path, result_path)\n",
    "\n",
    "# beams (ie beam width != 1). Write the formatted beam target and beam source files in a temp dir\n",
    "with tempfile.TemporaryDirectory() as temp_dir:\n",
    "  for beam_width in [3, 5, 10, 20, 50]:\n",
    "    beam_path = os.path.join(prediction_dir, 'beams_44000steps_' + str(beam_width) + '_test.npz')\n",
    "    processed_pred_path = os.path.join(temp_dir, 'pred_beam_' + str(beam_width))\n",
    "    processed_target_path = os.path.join(temp_dir, 'target_beam_' + str(beam_width))\n",
    "    processed_source_path = os.path.join(temp_dir, 'source_beam_' + str(beam_width))\n",
    "    result_path = os.path.join(cwd, 'beam_' + str(beam_width) + '_test_results')\n",
    "    \n",
    "    # reformat beam files\n",
    "    create_beam_source_target_parallel_txt_files(beam_path, vocab_path,\n",
    "                                               processed_pred_path,\n",
    "                                               target_path,\n",
    "                                               processed_target_path,\n",
    "                                               source_path,\n",
    "                                               processed_source_path, beam_width)\n",
    "    # evaluate reformatted beam predictions\n",
    "    evaluate_beam_by_rxn_class(processed_pred_path, processed_target_path, processed_source_path, result_path)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
